package com.qs.rocketmq.nacos.config.rule;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义负载均衡规则
 *
 * @author qinyupeng
 * @since 2022-03-25 10:45:49
 */
@Slf4j
public class MyTestRule extends AbstractLoadBalancerRule {

    private AtomicInteger index;

    public MyTestRule() {
        index = new AtomicInteger(0);
    }

    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {

    }

    @Override
    public Server choose(Object key) {
        ILoadBalancer loadBalancer = getLoadBalancer();
        List<Server> allServers = loadBalancer.getAllServers();
        if (CollectionUtils.isEmpty(allServers)) {
            log.error("注册中心未获取到可用的服务列表");
            return null;
        }

        for (int i = 0; i < allServers.size(); i++) {
            if (index.get() >= allServers.size()) {
                index.set(0);
            }

            Server server = allServers.get(index.getAndIncrement());
            if (server == null || !server.isAlive()) {
                continue;
            }

            log.info("负载均衡获取节点：host={}，port={}", server.getHost(), server.getHostPort());
            return server;
        }

        return null;
    }
}
